//
//  BRDF.hpp
//  RayTracer
//
//  Created by 俞云康 on 2/22/16.
//  Copyright © 2016 yushroo*m. All rights reserved.
//

#ifndef BRDF_hpp
#define BRDF_hpp

#include "Math.hpp"
#include "HitInfo.hpp"
#include "Sample.hpp"
#include "Fresnel.hpp"

enum BxDFType {
	BxDF_REFLECTION     = 1<<0,
	BxDF_TRANSMISSION   = 1<<1,
    BxDF_DIFFUSE        = 1<<2,
    BxDF_GLOSSY         = 1<<3,
    BxDF_SPECULAR       = 1<<4,
    BxDF_ALL_TYPES          = BxDF_DIFFUSE | BxDF_GLOSSY | BxDF_SPECULAR,
    BxDF_ALL_REFLECTION     = BxDF_REFLECTION | BxDF_ALL_TYPES,
    BxDF_ALL_TRANSMISSION   = BxDF_TRANSMISSION | BxDF_ALL_TYPES,
    BxDF_ALL                = BxDF_ALL_REFLECTION | BxDF_ALL_TRANSMISSION
};

class BxDF {
public:
	BxDFType type = BxDF_REFLECTION;
    const HitInfo* hit;
    
    BxDF(const HitInfo* hit, BxDFType type) : hit(hit), type(type) {}
    
    virtual ~BxDF() = default;
    
    bool matchesFlags(BxDFType flags) const {
        return (type & flags) == type;
    }
    
    virtual Color f(const Direction& woL, const Direction& wiL) const = 0;
    
    virtual Color sample_f(const Direction& woL, Direction* wiL, const Sample& sample, float*pdf) const {
        // Cosine-sample the hemisphere, flipping the direction if necessary
        *wiL = CosineSampleHemisphere(sample);
        if (woL.z < 0.f) wiL->z = -wiL->z;
        *pdf = this->pdf(woL, *wiL);
        return f(woL, *wiL);
    }
    
    // hemispherical-directional reflectance
    virtual Color rho(const Direction& w, const std::vector<Sample>& samples) const {
        // see PBRT p704
        Assert(!samples.empty());
        Color r(0, 0, 0);
        const size_t n_samples = samples.size();
        for (size_t i = 0; i < n_samples; ++i) {
            Direction wi;
            float pdf = 0.f;
            Color f = sample_f(w, &wi, samples[i], &pdf);
            if (pdf > 0.f)
                r += f * AbsCosTheta(wi);
        }
        return r / float(n_samples);
    }
    
    // hemispherical-hemispherical reflectance
    virtual Color rho(const std::vector<Sample>& samples_o, const std::vector<Sample>& samples_i) const {
        Assert((!samples_i.empty()) && (samples_i.size() == samples_o.size()));
        Color r(0, 0, 0);
        const size_t n_samples = samples_o.size();
        for (size_t i = 0; i < n_samples; ++i) {
            Direction wo, wi;
            auto& sa1 = samples_o[i];
            auto& sa2 = samples_i[i];
            wo = UniformSampleHemisphere(sa1.x, sa1.y);
            float pdf_o = INV_TWOPI, pdf_i = 0;
            Color f = sample_f(wo, &wi, sa2, &pdf_i);
            if (pdf_i > 0)
                r += f * AbsCosTheta(wi) * AbsCosTheta(wo) / (pdf_o * pdf_i);
        }
        return r / (M_PI * n_samples);
    }
    
    virtual float pdf(const Direction& woL, const Direction& wiL) const {
        return SameHemisphere(woL, wiL) ? AbsCosTheta(wiL) * INV_PI : 0;
    }
};


class SpecularReflection : public BxDF {
private:
    Color R;
    std::shared_ptr<Fresnel> fresnel;
public:
    SpecularReflection(const HitInfo* hit, const Color& r, std::shared_ptr<Fresnel> fresnel)
    : BxDF(hit, BxDFType(BxDF_REFLECTION | BxDF_SPECULAR)), R(r), fresnel(fresnel)
    {
    }
    
    virtual Color f(const Direction& woL, const Direction& wiL) const override {
        return Color(0, 0, 0);
    }
    
    virtual Color sample_f(const Direction& woL, Direction* wiL, const Sample& sample, float* pdf) const override;
};


class SpecularTransmission : public BxDF {
private:
    Color T;
    float eta_i, eta_t;
    FresnelDialectric fresnel;
    
public:
    SpecularTransmission(const HitInfo* hit, const Color& T, float eta_i, float eta_t)
    : BxDF(hit, BxDFType(BxDF_TRANSMISSION | BxDF_SPECULAR)), T(T), eta_i(eta_i), eta_t(eta_t), fresnel(eta_i, eta_t) {
	}
    
    virtual Color f(const Direction& wo, const Direction& wi) const override {
        return Color(0, 0, 0);
    }
    
    virtual Color sample_f(const Direction& wo, Direction* wi, const Sample& sample, float*pdf) const override;
};

#endif /* BRDF_hpp */
